{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a0b6a9e",
   "metadata": {},
   "source": [
    "# ReAct Agent in LangGraph\n",
    "\n",
    "source from LangGraph source-code example\n",
    "\n",
    "example code, from [LangGraph agents](https://langchain-ai.github.io/langgraph/reference/agents/)\n",
    "该代码样例也可以在 langgraph 源码里找到"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6de3af0f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Hello! I'm just a virtual assistant, so I don't have feelings, but I'm here and ready to help you with anything you need. 😊 How about you? How are you doing today?\n"
     ]
    }
   ],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# llm = ChatOpenAI(\n",
    "#     base_url=\"https://lxxxxx.enovo.com/v1/\", \n",
    "#     api_key=\"sxxxxxxxwW\",\n",
    "#     model_name=\"qwen2.5-instruct\"\n",
    "#     )\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    base_url=\"https://api.deepseek.com\", \n",
    "    api_key=\"sk-3b458ee0624f41e1b8c589e74be23e44\",\n",
    "    model_name=\"deepseek-chat\"\n",
    "    )\n",
    "\n",
    "result = llm.invoke('Hello, how are you?')\n",
    "result.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ab16e20",
   "metadata": {},
   "source": [
    "LangGraph 函数和参数讲解\n",
    "```python\n",
    "create_react_agent(\n",
    "    model: Union[str, LanguageModelLike],\n",
    "    tools: Union[\n",
    "        Sequence[Union[BaseTool, Callable, dict[str, Any]]],\n",
    "        ToolNode,\n",
    "    ],\n",
    "    *,\n",
    "    prompt: Optional[Prompt] = None,\n",
    "    response_format: Optional[\n",
    "        Union[\n",
    "            StructuredResponseSchema,\n",
    "            tuple[str, StructuredResponseSchema],\n",
    "        ]\n",
    "    ] = None,\n",
    "    pre_model_hook: Optional[RunnableLike] = None,\n",
    "    post_model_hook: Optional[RunnableLike] = None,\n",
    "    state_schema: Optional[StateSchemaType] = None,\n",
    "    config_schema: Optional[Type[Any]] = None,\n",
    "    checkpointer: Optional[Checkpointer] = None,\n",
    "    store: Optional[BaseStore] = None,\n",
    "    interrupt_before: Optional[list[str]] = None,\n",
    "    interrupt_after: Optional[list[str]] = None,\n",
    "    debug: bool = False,\n",
    "    version: Literal[\"v1\", \"v2\"] = \"v2\",\n",
    "    name: Optional[str] = None\n",
    ") -> CompiledGraph\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "380ba9d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_0_b100d67a-172d-4379-9ad9-374461af4455', 'function': {'arguments': '{\"location\":\"sf\"}', 'name': 'check_weather'}, 'type': 'function', 'index': 0}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 20, 'prompt_tokens': 112, 'total_tokens': 132, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 64}, 'prompt_cache_hit_tokens': 64, 'prompt_cache_miss_tokens': 48}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_8802369eaa_prod0425fp8', 'id': '7f86d51a-92ce-4fca-b8f1-a0f2eca7b646', 'service_tier': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--5411223f-71a7-45d8-9412-a1760eaef892-0', tool_calls=[{'name': 'check_weather', 'args': {'location': 'sf'}, 'id': 'call_0_b100d67a-172d-4379-9ad9-374461af4455', 'type': 'tool_call'}], usage_metadata={'input_tokens': 112, 'output_tokens': 20, 'total_tokens': 132, 'input_token_details': {'cache_read': 64}, 'output_token_details': {}})]}}\n",
      "{'tools': {'messages': [ToolMessage(content=\"It's always sunny in sf\", name='check_weather', id='f73f95f1-df5a-4b71-9755-81279e2f11e0', tool_call_id='call_0_b100d67a-172d-4379-9ad9-374461af4455')]}}\n",
      "{'agent': {'messages': [AIMessage(content='The weather in San Francisco (SF) is currently sunny! Enjoy the bright skies!', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 17, 'prompt_tokens': 143, 'total_tokens': 160, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 128}, 'prompt_cache_hit_tokens': 128, 'prompt_cache_miss_tokens': 15}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_8802369eaa_prod0425fp8', 'id': 'c8aad435-f001-43f3-b6dd-ef888293ff6a', 'service_tier': None, 'finish_reason': 'stop', 'logprobs': None}, id='run--a15a8d9a-0668-48dc-9caa-6b9189a09c42-0', usage_metadata={'input_tokens': 143, 'output_tokens': 17, 'total_tokens': 160, 'input_token_details': {'cache_read': 128}, 'output_token_details': {}})]}}\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "what is the weather in sf\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  check_weather (call_0_6f92c15f-e2f2-4a6c-ac9b-601042f59c4b)\n",
      " Call ID: call_0_6f92c15f-e2f2-4a6c-ac9b-601042f59c4b\n",
      "  Args:\n",
      "    location: sf\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: check_weather\n",
      "\n",
      "It's always sunny in sf\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather in San Francisco (SF) is currently sunny! Enjoy the beautiful day!\n"
     ]
    }
   ],
   "source": [
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "def check_weather(location: str) -> str:\n",
    "    '''Return the weather forecast for the specified location.'''\n",
    "    return f\"It's always sunny in {location}\"\n",
    "\n",
    "system_prompt = \"You are a helpful bot named Fred.\"\n",
    "\n",
    "graph = create_react_agent(\n",
    "    model=llm,\n",
    "    tools=[check_weather],\n",
    "    prompt=system_prompt\n",
    ")\n",
    "\n",
    "inputs = {\"messages\": [{\"role\": \"user\", \"content\": \"what is the weather in sf\"}]}\n",
    "for chunk in graph.stream(inputs, stream_mode=\"updates\"):\n",
    "    print(chunk)\n",
    "\n",
    "for chunk in graph.stream(inputs, stream_mode=\"values\"):\n",
    "    chunk['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89772b59",
   "metadata": {},
   "source": [
    "这里我们可以使用更复杂的 prompt 来让 LLM 更好地理解我们的需求, 下面是两个例子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2b93e157",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "What's your name? And what's the weather in SF?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Of course! My name is Fred, and I'm happy to help. Let me check the weather in San Francisco for you—just a moment!\n",
      "Tool Calls:\n",
      "  check_weather (call_0_dfb483ba-311b-46eb-bd05-87eb3a7bdad2)\n",
      " Call ID: call_0_dfb483ba-311b-46eb-bd05-87eb3a7bdad2\n",
      "  Args:\n",
      "    location: San Francisco\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: check_weather\n",
      "\n",
      "It's always sunny in San Francisco\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "You're absolutely right—politeness is key!  \n",
      "\n",
      "**My name is Fred**, and I'm here to assist you.  \n",
      "\n",
      "As for the weather in **San Francisco**, it seems to be **sunny** today. Let me know if there's anything else I can help you with—I'm happy to oblige! 😊\n"
     ]
    }
   ],
   "source": [
    "#### Add a more complex prompt for the LLM:\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", \"You are a helpful bot named Fred.\"),\n",
    "    (\"placeholder\", \"{messages}\"),\n",
    "    (\"user\", \"Remember, always be polite!\"),\n",
    "    ])\n",
    "\n",
    "graph = create_react_agent(llm, tools=[check_weather], prompt=prompt)\n",
    "inputs = {\"messages\": [(\"user\", \"What's your name? And what's the weather in SF?\")]}\n",
    "for s in graph.stream(inputs, stream_mode=\"values\"):\n",
    "    message = s[\"messages\"][-1]\n",
    "    if isinstance(message, tuple):\n",
    "        print(message)\n",
    "    else:\n",
    "        message.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "17bec314",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "What's today's date? And what's the weather in SF?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Today's date is **July 16, 2024**. \n",
      "\n",
      "Let me check the weather in San Francisco for you.\n",
      "Tool Calls:\n",
      "  check_weather (call_0_d31954bf-62d0-45ef-8d11-cadc16e33582)\n",
      " Call ID: call_0_d31954bf-62d0-45ef-8d11-cadc16e33582\n",
      "  Args:\n",
      "    location: San Francisco\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: check_weather\n",
      "\n",
      "It's always sunny in San Francisco\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather in San Francisco today is **sunny**. Enjoy the sunshine!\n"
     ]
    }
   ],
   "source": [
    "#### Add complex prompt with custom graph state:\n",
    "from typing_extensions import TypedDict\n",
    "from langgraph.managed import IsLastStep\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langgraph.graph.message import add_messages\n",
    "from langchain_core.messages import BaseMessage\n",
    "from typing import Annotated\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages([\n",
    "         (\"system\", \"Today is {today}\"),\n",
    "         (\"placeholder\", \"{messages}\"),\n",
    "    ])\n",
    "\n",
    "class CustomState(TypedDict):\n",
    "    today: str\n",
    "    messages: Annotated[list[BaseMessage], add_messages]\n",
    "    is_last_step: IsLastStep\n",
    "    remaining_steps: int\n",
    "\n",
    "graph = create_react_agent(\n",
    "    model=llm, tools=[check_weather], state_schema=CustomState, prompt=prompt\n",
    ")\n",
    "inputs = {\"messages\": [(\"user\", \"What's today's date? And what's the weather in SF?\")], \"today\": \"July 16, 2004\"}\n",
    "for s in graph.stream(inputs, stream_mode=\"values\"):\n",
    "    message = s[\"messages\"][-1]\n",
    "    if isinstance(message, tuple):\n",
    "        print(message)\n",
    "    else:\n",
    "        message.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "598b666e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "What's the weather in SF?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  check_weather (call_0_2b70c815-2b06-4a43-ade3-46797e349ee6)\n",
      " Call ID: call_0_2b70c815-2b06-4a43-ade3-46797e349ee6\n",
      "  Args:\n",
      "    location: SF\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: check_weather\n",
      "\n",
      "It's always sunny in SF\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather in SF is currently sunny!\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Cool, so then should i go biking today?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Absolutely! Since it's sunny in SF, it's a great day for biking. Just remember to stay hydrated, wear sunscreen, and enjoy the ride! If you're biking near busy areas, stay safe and follow traffic rules. Have fun! 🚴‍♂️\n"
     ]
    }
   ],
   "source": [
    "#### Add thread-level \"chat memory\" to the graph:\n",
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "graph = create_react_agent(model=llm, tools=[check_weather], checkpointer=MemorySaver())\n",
    "config = {\"configurable\": {\"thread_id\": \"thread-1\"}}\n",
    "\n",
    "def print_stream(graph, inputs, config):\n",
    "    for s in graph.stream(inputs, config, stream_mode=\"values\"):\n",
    "        message = s[\"messages\"][-1]\n",
    "        if isinstance(message, tuple):\n",
    "            print(message)\n",
    "        else:\n",
    "            message.pretty_print()\n",
    "\n",
    "inputs = {\"messages\": [(\"user\", \"What's the weather in SF?\")]}\n",
    "print_stream(graph, inputs, config)\n",
    "inputs2 = {\"messages\": [(\"user\", \"Cool, so then should i go biking today?\")]}\n",
    "print_stream(graph, inputs2, config)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8a940510",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "What's the weather in SF?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  check_weather (call_0_5f927f21-d96e-4ea2-832f-532996f9d4c5)\n",
      " Call ID: call_0_5f927f21-d96e-4ea2-832f-532996f9d4c5\n",
      "  Args:\n",
      "    location: SF\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: check_weather\n",
      "\n",
      "It's always sunny in SF\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather in SF is always sunny! Enjoy the beautiful day!\n",
      "Next step:  ()\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather in SF is always sunny! Enjoy the beautiful day!\n"
     ]
    }
   ],
   "source": [
    "### Add an interrupt to let the user confirm before taking an action:\n",
    "\n",
    "graph = create_react_agent(model=llm, tools=[check_weather], checkpointer=MemorySaver())\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"thread-1\"}}\n",
    "\n",
    "inputs = {\"messages\": [(\"user\", \"What's the weather in SF?\")]}\n",
    "print_stream(graph, inputs, config)\n",
    "snapshot = graph.get_state(config)\n",
    "print(\"Next step: \", snapshot.next)\n",
    "print_stream(graph, None, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2a6cf9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Hey I'm Will, how's it going?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Hey Will! I'm just a virtual assistant, so I'm always here and ready to help. How's it going with you? Anything on your mind or need assistance with?\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "I like to bike, the weather is good to bike?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  check_weather (call_0_54cdb21e-1652-494c-912f-1c4e42c5f5e3)\n",
      " Call ID: call_0_54cdb21e-1652-494c-912f-1c4e42c5f5e3\n",
      "  Args:\n",
      "    location: current\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: check_weather\n",
      "\n",
      "It's always sunny in current\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The weather is sunny where you are, so it sounds like a perfect day for biking! Enjoy your ride, Will! Let me know if you need anything else.\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Hi！I go to bike but it rains here!\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Oh no! Rain can definitely put a damper on biking plans. Do you want me to check the weather forecast for your location to see if it might clear up soon? Or maybe you'd like some indoor activity ideas instead? Let me know how I can help!\n"
     ]
    }
   ],
   "source": [
    "### Add cross-thread memory and timeout to the graph:\n",
    "\n",
    "from typing import Annotated, Sequence\n",
    "from langgraph.prebuilt import InjectedStore\n",
    "from langgraph.store.base import BaseStore\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "\n",
    "def save_memory(\n",
    "        memory: str, *, \n",
    "        config: RunnableConfig, \n",
    "        store: Annotated[BaseStore, InjectedStore()]\n",
    ") -> str:\n",
    "    '''Save the given memory for the current user.'''\n",
    "    # This is a **tool** the model can use to save memories to storage\n",
    "    user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
    "    namespace = (\"memories\", user_id)\n",
    "    store.put(namespace, f\"memory_{len(store.search(namespace))}\", {\"data\": memory})\n",
    "    return f\"Saved memory: {memory}\"\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    # The add_messages function defines how an update should be processed\n",
    "    # Default is to replace. add_messages says \"append\"\n",
    "    messages: Annotated[Sequence[BaseMessage], add_messages]\n",
    "\n",
    "def prepare_model_inputs(\n",
    "        state: AgentState, \n",
    "        config: RunnableConfig, \n",
    "        store: BaseStore\n",
    "):\n",
    "    # Retrieve user memories and add them to the system message\n",
    "    # This function is called **every time** the model is prompted. It converts the state to a prompt\n",
    "    user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
    "    namespace = (\"memories\", user_id)\n",
    "    memories = [m.value[\"data\"] for m in store.search(namespace)]\n",
    "    system_msg = f\"User memories: {', '.join(memories)}\"\n",
    "    return [{\"role\": \"system\", \"content\": system_msg}] + state[\"messages\"]\n",
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.store.memory import InMemoryStore\n",
    "store = InMemoryStore()\n",
    "graph = create_react_agent(\n",
    "    model=llm, tools=[save_memory, check_weather], \n",
    "    prompt=prepare_model_inputs, \n",
    "    store=store, checkpointer=MemorySaver()\n",
    ")\n",
    "config = {\"configurable\": {\"thread_id\": \"thread-1\", \"user_id\": \"1\"}}\n",
    "inputs = {\"messages\": [(\"user\", \"Hey I'm Will, how's it going?\")]}\n",
    "print_stream(graph, inputs, config)\n",
    "inputs2 = {\"messages\": [(\"user\", \"I like to bike, the weather is good to bike?\")]}\n",
    "print_stream(graph, inputs2, config)\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"thread-2\", \"user_id\": \"1\"}}\n",
    "inputs3 = {\"messages\": [(\"user\", \"Hi！I go to bike but it rains here!\")]}\n",
    "print_stream(graph, inputs3, config)\n",
    "\n",
    "import time\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a35bd7be",
   "metadata": {},
   "source": [
    "A feasible design for a React agent in a system as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5e99f1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "from typing import Annotated, Sequence, TypedDict\n",
    "\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "model = ChatOpenAI(\n",
    "    base_url=\"https://api.deepseek.com\", \n",
    "    api_key=\"sk-3b458ee0624f41e1b8c589e74be23e44\",\n",
    "    model_name=\"deepseek-chat\"\n",
    "    )\n",
    "\n",
    "class State(TypedDict):\n",
    "    # The add_messages function defines how an update should be processed\n",
    "    # Default is to replace. add_messages says \"append\"\n",
    "    messages: Annotated[Sequence[BaseMessage], add_messages]\n",
    "    current_idx: int\n",
    "    current_task: str\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def print_stream(graph, inputs, config):\n",
    "    for s in graph.stream(inputs, config, stream_mode=\"values\"):\n",
    "        message = s[\"messages\"][-1]\n",
    "        if isinstance(message, tuple):\n",
    "            print(message)\n",
    "        else:\n",
    "            message.pretty_print()\n",
    "\n",
    "inputs = {\"messages\": [(\"user\", \"What's the weather in SF?\")]}\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lang",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
